# 指令替换
# obf2
from importlib.metadata import entry_points
import string
from typing import Set
import idc
import idaapi
import ida_kernwin
import ida_funcs
import idautils

from example.ida.utils import guess_machine,expr2colorstr
from future.utils import viewitems, viewvalues

from miasm.arch.x86.regs import *
from miasm.core.bin_stream_ida import bin_stream_ida
from miasm.expression.simplifications import expr_simp
from miasm.ir.ir import IRBlock, AssignBlock
from miasm.ir.symbexec import SymbolicExecutionEngine
from miasm.analysis.data_flow import DeadRemoval, merge_blocks, remove_empty_assignblks
from miasm.analysis.cst_propag import propagate_cst_expr
from miasm.analysis.data_flow import load_from_int
from miasm.analysis.data_flow import ReachingDefinitions,DiGraphDefUse
from miasm.expression.expression import ExprLoc, ExprInt, ExprOp, ExprAssign,ExprCond, ExprId, ExprMem,ExprSlice
from miasm.analysis.simplifier import IRCFGSimplifierCommon, IRCFGSimplifierSSA
from miasm.core.locationdb import LocationDB
from miasm.core.asmblock import AsmCFG,AsmBlock,AsmConstraint
from miasm.ir.translators.translator import Translator
from miasm.ir.ir import IRCFG
from draw import *
from BaseSolver import BaseSolver
from mylib import ActionList
# 空间地址上的下一块为next
def asmblock_del_next(action:ActionList,asmcfg:AsmCFG,asmblk:AsmBlock):
    del_list=[]
    for bto in asmblk.bto:
        if bto.c_t=='c_next':
            del_list.append(bto.loc_key)

    for dst_loc in del_list:
        asmcfg.del_edge(asmblk.loc_key,dst_loc)
        action.graph_del_edge(asmblk.get_offsets()[0],asmcfg.loc_key_to_block(dst_loc).get_offsets()[0])
        

# 因为跳转等指令的下一块为to
def asmblock_del_to(action:ActionList,asmcfg:AsmCFG,asmblk:AsmBlock):
    del_list=[]
    for bto in asmblk.bto:
        if bto.c_t=='c_to':
            del_list.append(bto.loc_key)

    for dst_loc in del_list:
        asmcfg.del_edge(asmblk.loc_key,dst_loc)
        action.graph_del_edge(asmblk.get_offsets()[0],asmcfg.loc_key_to_block(dst_loc).get_offsets()[0])
def get_asmcfg_addr(offset,mdis):
    asmcfg_more=mdis.dis_multiblock(offset)
    return asmcfg_more

def merge_graph(origin_graph,new_graph):
    # Add known blocks
    for block in new_graph.blocks:
        origin_graph.add_block(block)
    # Add nodes not already in it (ie. not linked to a block)
    for node in new_graph.nodes():
        origin_graph.add_node(node)
    # -> add_edge(x, y, constraint)
    for edge in new_graph._edges:
        # May fail if there is an incompatibility in edges constraints
        # between the two graphs
        print(edge)
        origin_graph.add_edge(*edge, constraint=new_graph.edges2constraint[edge])
    return
# 扫描每个汇编块
scanner_blocks=[]
def scanner_asmblock(action:ActionList,asmcfg:AsmCFG,asmblk_loc_key:AsmBlock.loc_key,mdis):
    global scanner_blocks
    # 最后是否为JZ/JNZ引起的跳转
    asmblk=asmcfg.loc_key_to_block(asmblk_loc_key)
    
    if asmblk is None:
        return
    last_instr=asmblk.lines[-1]
    last_instr_name=last_instr.name
    print('scanning',hex(asmblk.get_offsets()[0]))
    # print('last',last_instr_name)
    if last_instr_name=='JZ':
        jz_dst=last_instr.args[0].loc_key
        # 判断下一块是否为jnz
        next_blk_key=asmblk.get_next()
        next_blk=asmcfg.loc_key_to_block(next_blk_key)
        next_last_instr=next_blk.lines[-1]
        # print('next',next_last_instr.name,next_last_instr.args[0].loc_key,jz_dst)
        if next_last_instr.name=='JNZ' and next_last_instr.args[0].loc_key==jz_dst:
            # 不透明跳转
            # jz的上一个push为返回地址
            push_instr=asmblk.lines[-2]
            if push_instr.name=='PUSH':
                # 压入返回地址
                ret_addr=int(push_instr.args[0])
                ret_blk=asmcfg.getby_offset(ret_addr)
                if ret_blk is None:
                    print('merge',hex(ret_addr))
                    asmcfg_more=get_asmcfg_addr(ret_addr,mdis)
                    if asmcfg_more is not None:
                        asmcfg.merge(asmcfg_more)
                        for blks in asmcfg_more.blocks:
                            scanner_blocks.append(blks)

                        ret_blk=asmcfg.getby_offset(ret_addr)
                    else:
                        return
                print('----------addr',ret_addr)
                # 得到next块
                action.graph_del_blk(next_blk.get_offsets()[0])
                asmcfg.del_block(next_blk)
                asmblock_del_next(action,asmcfg,asmblk)
                asmblock_del_to(action,asmcfg,asmblk)
                asmcfg.add_edge(asmblk.loc_key,ret_blk.loc_key,AsmConstraint(ret_blk.loc_key,'c_next'))
                action.graph_add_edge(asmblk.get_offsets()[0],ret_blk.get_offsets()[0])
                asmblk.lines[-1].name='CALL'
                # asmblk.lines[-1].args=[next_last_instr.args[0]]
                jmp_offset=asmcfg.loc_key_to_block(jz_dst).get_offsets()[0]
                asmblk.lines[-1].args=[ExprInt(jmp_offset,32)]
        
    
    # print(last_instr_name,last_instr_args)
    return
# 删去不重要的分支
def cut_cfg(action:ActionList,asmcfg:AsmCFG,start_addr):
    del_blocks=[]
    for blk in asmcfg.blocks:
        pred_counts=asmcfg.predecessors(blk.loc_key)
        succ_counts=asmcfg.successors(blk.loc_key)
        if blk.get_offsets()[0]<start_addr:
            del_blocks.append(blk)
            continue
        if pred_counts==0 and succ_counts==0:
            del_blocks.append(blk)
            continue
        
    for blk in del_blocks:
        action.graph_del_blk(blk.get_offsets()[0])
        asmcfg.del_block(blk)
    return
def sub_maze(addr,machine,asmcfg:AsmCFG,head,mdis,loc_db):
    action=ActionList()
    func_addr=addr
    draw_graph(asmcfg,'before')
    asmcfg_cut=asmcfg.copy()
    # ircfg=lifter.new_ircfg_from_asmcfg(asmcfg)
    # 得到入口Loc_key
    entry_points=set([loc_db.get_offset_location(func_addr)])
    # 流程图的头节点
    head=list(entry_points)[0]
    head_blk=asmcfg.loc_key_to_block(head)
    head_off=head_blk.get_offsets()[0]
    
    global scanner_blocks
    for blks in asmcfg.blocks:
        scanner_blocks.append(blks)
    # 遍历asmcfg
    while True:
        
        blks_len=len(scanner_blocks)    
        for asmblock in scanner_blocks:
            scanner_asmblock(action,asmcfg_cut,asmblock.loc_key,mdis)
            bad_list=[]
            for bad in asmcfg_cut.get_bad_blocks():
                # print('bad',bad)
                bad_list.append(bad)
            for bad in bad_list:
                asmcfg_cut.del_block(bad)
        while blks_len>0:
            del scanner_blocks[0]
            blks_len-=1
        if len(scanner_blocks)<=0:
            break
        
    cut_cfg(action,asmcfg_cut,func_addr)
    asmcfg_cut.get_bad_blocks()
    draw_graph(asmcfg_cut,'asmcfg')
    return action
class MazeSolver(BaseSolver):
    def __init__(self) -> None:
        super().__init__('Maze')
        pass
    
    def Solve(self,obj):
        super().Solve(obj)
        self.action=sub_maze(self.addr,self.machine,self.asmcfg,self.head,self.mdis,self.loc_db)
        return 